import path from 'path'
import db from '../db/inedx.js'
import { saveJson, sleep, random, readJson, queue } from '../../../utils/util.js'
import { rootDir, techniqueDir, subTechniqueDir } from '../config/index.js'

import { glob } from 'glob'

import OpenAI from './llm/open-ai.js'
import NM from './llm/nm.js'

class Translate {
  constructor() {
    this.api = new OpenAI()
    // this.api = new NM()
  }

  get db() {
    return db.projectConfig
  }

  get config() {
    return db.projectConfig.data
  }

  //  翻译攻击技术
  async translateTechnique() {
    let files = await glob(`${techniqueDir}/*.json`)

    await queue(
      files,
      async (file, index) => {
        // 已经执行
        if (index < this.config.translateTechniqueIndex) return

        let technique = await readJson(file)

        // 400输入数据可能包含不适当的内容
        if (['T1559', 'T1560'].includes(technique.id)) return

        //  进度
        console.log(`当前进度:${index + 1}/${files.length} ${technique.id}`)

        delete technique.descriptionHtml
        technique.description = this.filterBracketedNumbers(technique.description)

        // 全文翻译
        // technique = await this.api.translateFile(technique)

        const error = e => {
          console.error('翻译失败', e)
        }

        await Promise.all([
          this.translateObject(technique, ['title', 'description']).catch(error),
          this.translateObject(technique.details, ['tactic', 'created', 'lastModified']).catch(error),
          this.translateObjectArray(technique.subTechniques, ['name']).catch(error),
          this.translateObjectArray(technique.mitigations, ['description']).catch(error),
          this.translateObjectArray(technique.procedureExamples, ['description']).catch(error),
          this.translateObjectArray(technique.detections, ['dataComponent']).catch(error),
          this.translateObjectArray(technique.detections, ['description']).catch(error),
          this.translateObjectArray(technique.references, ['text']).catch(error),
        ])

        saveJson(path.join(file.replace(/technique/g, 'translate/technique')), technique)
      },
      {
        limit: 1,
        onStage: async index => {
          if (index > this.config.translateTechniqueIndex) {
            this.config.translateTechniqueIndex = index
            await this.db.write()
          }
        },
      }
    )
  }

  //  翻译子攻击技术
  async subTranslateTechnique() {
    let files = await glob(`${subTechniqueDir}/*.json`)

    await queue(
      files,
      async (file, index) => {
        // 已经执行
        if (index < this.config.translateSubTechniqueIndex) return

        let subTechnique = await readJson(file)

        //  进度
        console.log(`当前进度:${index + 1}/${files.length} ${subTechnique.id}`)

        delete subTechnique.descriptionHtml
        subTechnique.description = this.filterBracketedNumbers(subTechnique.description)

        // 全文翻译
        // subTechnique = await this.api.translateFile(subTechnique)

        const error = (e, data) => console.error('翻译失败', e, data)

        await Promise.all([
          this.translateObject(subTechnique, ['title', 'description']).catch(err => error(err, '1')),
          this.translateObject(subTechnique.details, ['tactic', 'created', 'lastModified']).catch(err =>
            error(err, '2')
          ),
          this.translateObjectArray(subTechnique.subTechniques, ['name']).catch(err => error(err, '3')),
          this.translateObjectArray(subTechnique.mitigations, ['description']).catch(err => error(err, '4')),
          this.translateObjectArray(subTechnique.procedureExamples, ['description']).catch(err =>
            error(err, '5')
          ),
          this.translateObjectArray(subTechnique.detections, ['dataComponent']).catch(err => error(err, '6')),
          this.translateObjectArray(subTechnique.detections, ['description']).catch(err => error(err, '7')),
          this.translateObjectArray(subTechnique.references, ['text']).catch(err => error(err, '8')),
        ])

        saveJson(path.join(file.replace(/sub-technique/g, 'translate/sub-technique')), subTechnique)
      },
      {
        limit: 1,
        onStage: async index => {
          if (index > this.config.translateSubTechniqueIndex) {
            this.config.translateSubTechniqueIndex = index
            await this.db.write()
          }
        },
      }
    )
  }

  // 翻译对象
  async translateObject(data, keys) {
    if (Object.keys(data).length == 0) return

    const texts = keys.map(key => data[key])

    const results = await this.api.translateTexts(texts)

    keys.forEach((key, index) => {
      data[key] = results[index]
    })
  }

  // 翻译对象数组
  async translateObjectArray(array, keys) {
    if (!array || array?.length == 0) return

    if (keys.length == 1) {
      const key = keys[0]
      const texts = []

      array.forEach(item => {
        texts.push(item[key])
      })

      const results = await this.api.translateTexts(texts)

      array.forEach((item, index) => {
        try {
          item[key] = results[index]
        } catch (e) {
          console.error('翻译失败', {
            texts,
            results,
            array,
            keys,
          })
          throw e
        }
      })
    } else {
      const tasks = []

      array.forEach(item => {
        tasks.push(async () => {
          await this.translateObject(item, keys)
        })
      })

      await Promise.all(tasks)
    }
  }

  filterBracketedNumbers(text) {
    // 使用正则表达式匹配 [数字] 并替换为空字符串
    return text.replace(/\[\d+\]/g, '')
  }
}

export default Translate
